from __main__ import *
import os
import subprocess
import logging
import re
from utils import find_files

CMD_REGEX = re.compile(r"([0-9a-fA-F]+):\s+([0-9a-fA-F]+)\s+([0-9a-fA-F]+)\s+(.*)")

class Disassembly:
    def __init__(self, compilerPath: str, elfFile: str):
        logging.info(f"Running {__class__}")
        # find objdump
        self.__objdump = find_files(compilerPath, "objdump")
        if self.__objdump is None:
            logging.error(f"The Cross Compiler is {compilerPath} is invalid")
            raise FileNotFoundError("objdump")
        logging.debug("Found fromelf command. Path is: " + self.__objdump)

        # check elf file exists
        if not os.path.exists(elfFile):
            logging.error("Unable to locate ELF file. Please check filepath entered.")
            raise FileNotFoundError(elfFile)
        self.__elf = elfFile
        self.__disassembleOutput = os.path.splitext(elfFile)[0] + "_asm.s"
        logging.debug("Found ELF file. Path is: " + self.__elf)

        self.__disassemble()

    def __disassemble(self):
        if os.path.exists(self.__disassembleOutput):
            logging.info("Disassembly file already exists. Skipping disassembly.")
            return
        # Actually run the fromelf command
        try:
            command = (
                f'"{self.__objdump}" -d "{self.__elf}" > "{self.__disassembleOutput}"'
            )
            output = subprocess.check_output(
                command, shell=True, universal_newlines=True
            )
        except subprocess.CalledProcessError as e:
            logging.error(
                f"Failed to run the fromelf command: {command}. Please try running the command manually."
            )
            raise e
        logging.info("SUCCESSFULLY produced disassembly file using Arm Compiler")

    def getDisassemble(self):
        return self.__disassembleOutput


def DisassemblyToArrayGNU(filename: str) -> list[list[str]]:
    disassembly_array: list[list[str]] = []
    with open(filename) as f:
        content = f.readlines()
        address_pattern = re.compile("[0-9a-fA-F]")
        for line in content:
            line = line.strip()
            try:
                address, instruction = line.split(":\t")
            except:
                continue
            if address_pattern.match(address):
                disassembly_array.append([address, instruction.split(" \t")[1]])

    return disassembly_array


def DisassemblyToArrayFROMELF(filename: str) -> list[str]:
    disassembly_array: list[str] = []
    with open(filename) as f:
        content = f.readlines()
        address_pattern = re.compile("^0x[0-9a-fA-F]{8}$")
        for line in content:
            line = line.strip()
            if line.startswith("0x"):
                columns = line.split("    ", 3)
                if len(columns) >= 3:
                    address = columns[0].strip(":")
                    instruction = columns[3]
                    if address_pattern.match(address):
                        disassembly_array.append([address.lstrip("0x"), instruction])

    return disassembly_array


def DisassemblyToDict(list):
    # converting Disassembly list into dictionary
    lst = [i[0] for i in list]
    dTrans = {k: v for v, k in enumerate(lst)}
    return dTrans


class DisassemblyDecoder:
    def __init__(self, disassemblyFile: str):
        self.__file = disassemblyFile
        self.__aDisassembly = DisassemblyToArrayGNU(disassemblyFile)
        self.__dTrans = DisassemblyToDict(self.__aDisassembly)

    def get(self, saddr: str, eaddr:str):
        try:
            start = self.__dTrans[saddr]
            end = self.__dTrans[eaddr]
        except:
            return None
        return self.__aDisassembly[start : end]


if __name__ == "__main__":
    d = DisassemblyDecoder("mydump/bindir/vmlinux_asm.s")
    print(d.get("0x00000000"))